"""
rembg 背景去除封装
"""

from __future__ import annotations

import io
from typing import Optional

import numpy as np
from PIL import Image


class RembgError(Exception):
    """rembg 处理异常"""


class RembgRunner:
    """
    rembg 背景去除执行器
    """

    def __init__(self):
        try:
            from rembg import remove  # type: ignore
        except ImportError as exc:
            raise RembgError("未检测到 rembg 库，请先安装。") from exc

        self._remove = remove

    def remove_background(
        self,
        image: np.ndarray,
        alpha_matting: bool = False,
        alpha_matting_foreground_threshold: int = 240,
        alpha_matting_background_threshold: int = 10,
        alpha_matting_erode_structure_size: int = 10,
        alpha_matting_base_size: int = 1000
    ) -> np.ndarray:
        """
        对图片执行背景去除

        Args:
            image: 输入图片 (支持RGB/RGBA/灰度)
            alpha_matting: 是否开启alpha_matting
            alpha_matting_foreground_threshold: 前景阈值
            alpha_matting_background_threshold: 背景阈值
            alpha_matting_erode_structure_size: 侵蚀结构尺寸
            alpha_matting_base_size: 基准尺寸

        Returns:
            np.ndarray: 去除背景后的 RGBA 图片
        """
        if image is None:
            raise RembgError("输入图片为空")

        if image.ndim not in (2, 3):
            raise RembgError(f"不支持的图片维度: {image.ndim}")

        if image.ndim == 2:
            image = np.stack([image] * 3, axis=-1)

        if image.shape[2] == 4:
            pil_image = Image.fromarray(image.astype(np.uint8), mode="RGBA")
        else:
            pil_image = Image.fromarray(image[:, :, :3].astype(np.uint8), mode="RGB")

        try:
            output = self._remove(
                pil_image,
                alpha_matting=alpha_matting,
                alpha_matting_foreground_threshold=alpha_matting_foreground_threshold,
                alpha_matting_background_threshold=alpha_matting_background_threshold,
                alpha_matting_erode_structure_size=alpha_matting_erode_structure_size,
                alpha_matting_base_size=alpha_matting_base_size
            )
        except Exception as exc:
            raise RembgError(f"rembg 处理失败：{exc}") from exc

        result_image = self._to_image(output)
        return np.array(result_image.convert("RGBA"), dtype=np.uint8)

    @staticmethod
    def _to_image(result: Optional[object]) -> Image.Image:
        """
        将 rembg 输出转换为 PIL Image

        Args:
            result: rembg 输出

        Returns:
            Image.Image: 转换后的图片
        """
        if result is None:
            raise RembgError("rembg 返回结果为空")

        if isinstance(result, Image.Image):
            return result

        if isinstance(result, bytes):
            return Image.open(io.BytesIO(result))

        raise RembgError(f"无法解析 rembg 输出类型: {type(result)}")

